#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
财联社电报数据模型 - PostgreSQL版本
"""

import sys
import time
import os
from datetime import datetime
from core.config import PG_HOST, PG_PORT, PG_DB_NAME as PG_DATABASE, PG_USER, PG_PASSWORD

# 添加Web配置导入路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from web.web_config import web_config

# 尝试导入psycopg2，如果失败则设置标志
try:
    import psycopg2
    from psycopg2 import OperationalError
    PSYCOPG2_AVAILABLE = True
except ImportError:
    print("警告: 无法导入psycopg2模块，PostgreSQL功能将不可用")
    PSYCOPG2_AVAILABLE = False


class PostgresTelegramDatabase:
    def __init__(self):
        # 使用配置文件中的PostgreSQL连接参数
        self.host = PG_HOST
        self.port = PG_PORT
        self.database = PG_DATABASE
        self.user = PG_USER
        self.password = PG_PASSWORD
        self.init_database()

    def get_connection(self):
        """获取数据库连接"""
        if not PSYCOPG2_AVAILABLE:
            print("错误: psycopg2模块不可用，无法连接到PostgreSQL数据库")
            return None
        
        try:
            conn = psycopg2.connect(
                host=self.host,
                port=self.port,
                database=self.database,
                user=self.user,
                password=self.password
            )
            return conn
        except OperationalError as e:
            print(f"连接PostgreSQL数据库时出错: {e}")
            # 如果连接失败，尝试创建数据库
            if "database \"{}\" does not exist".format(self.database) in str(e):
                self._create_database()
                return self.get_connection()
            return None
        except Exception as e:
            print(f"数据库连接异常: {e}")
            return None

    def _create_database(self):
        """创建数据库"""
        try:
            # 首先连接到默认数据库
            conn = psycopg2.connect(
                host=self.host,
                port=self.port,
                database="postgres",  # 默认数据库
                user=self.user,
                password=self.password
            )
            conn.autocommit = True
            cursor = conn.cursor()
            cursor.execute(f"CREATE DATABASE {self.database}")
            print(f"数据库 {self.database} 创建成功")
            cursor.close()
            conn.close()
        except Exception as e:
            print(f"创建数据库时出错: {e}")

    def init_database(self):
        """初始化数据库表"""
        conn = self.get_connection()
        if not conn:
            print("无法连接到数据库，初始化失败")
            return

        try:
            cursor = conn.cursor()
            
            # 创建电报表 - 添加字段注释
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS telegrams (
                    id SERIAL PRIMARY KEY,                  -- 自增主键
                    telegram_id INTEGER,                    -- 电报ID
                    time TEXT,                              -- 发布时间
                    title TEXT,                             -- 标题
                    content TEXT,                           -- 内容
                    url TEXT,                               -- 网页链接
                    mobile_url TEXT,                        -- 移动端链接
                    reading_num INTEGER,                    -- 阅读数
                    id_original INTEGER UNIQUE,             -- 原始ID（唯一标识符）
                    audio_url TEXT,                         -- 音频链接
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, -- 创建时间
                    is_pushed BOOLEAN DEFAULT FALSE         -- 是否已推送
                )
            ''')
            
            # 创建索引提高查询效率
            cursor.execute('''
                CREATE INDEX IF NOT EXISTS idx_telegram_id ON telegrams(telegram_id)
            ''')
            
            cursor.execute('''
                CREATE INDEX IF NOT EXISTS idx_time ON telegrams(time)
            ''')
            
            cursor.execute('''
                CREATE INDEX IF NOT EXISTS idx_id_original ON telegrams(id_original)
            ''')
            
            conn.commit()
            print("PostgreSQL数据库表初始化成功")
        except Exception as e:
            print(f"初始化数据库表时出错: {e}")
            conn.rollback()
        finally:
            cursor.close()
            conn.close()

    def save_telegram(self, telegram_data):
        """保存单条电报数据"""
        conn = self.get_connection()
        if not conn:
            return False

        try:
            cursor = conn.cursor()
            
            # 使用INSERT ... ON CONFLICT DO UPDATE，相当于SQLite的INSERT OR REPLACE
            cursor.execute('''
                -- 获取Web配置中的推送开关状态
                -- 如果未开启推送，则默认标记为已推送；如果开启推送，则默认标记为未推送
                INSERT INTO telegrams 
                (telegram_id, time, title, content, url, mobile_url, reading_num, id_original, audio_url, is_pushed)
                VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                ON CONFLICT (id_original) DO UPDATE 
                SET telegram_id = EXCLUDED.telegram_id,
                    time = EXCLUDED.time,
                    title = EXCLUDED.title,
                    content = EXCLUDED.content,
                    url = EXCLUDED.url,
                    mobile_url = EXCLUDED.mobile_url,
                    reading_num = EXCLUDED.reading_num,
                    audio_url = EXCLUDED.audio_url;
                    -- 不再保持原有推送状态不变，确保已推送的消息保持已推送状态
              ''', (
                telegram_data.get('id'),
                telegram_data.get('time'),
                telegram_data.get('title'),
                telegram_data.get('content'),
                telegram_data.get('url'),
                telegram_data.get('mobileUrl'),
                telegram_data.get('reading_num'),
                telegram_data.get('id_original'),  # 这是真正的唯一标识符
                telegram_data.get('audio_url'),
                not web_config.get_wechat_push_enabled()  # 根据推送开关设置默认推送状态
            ))
            
            conn.commit()
            return cursor.rowcount > 0  # 返回是否插入/更新了数据
        except Exception as e:
            print(f"保存电报数据时出错: {e}")
            conn.rollback()
            return False
        finally:
            cursor.close()
            conn.close()

    def save_telegrams(self, telegrams_data):
        """批量保存电报数据"""
        conn = self.get_connection()
        if not conn:
            return 0

        updated_count = 0
        try:
            cursor = conn.cursor()
            
            for telegram in telegrams_data:
                try:
                    # 使用INSERT ... ON CONFLICT DO UPDATE
                    cursor.execute('''
                        -- 获取Web配置中的推送开关状态
                        -- 如果未开启推送，则默认标记为已推送；如果开启推送，则默认标记为未推送
                        INSERT INTO telegrams 
                        (telegram_id, time, title, content, url, mobile_url, reading_num, id_original, audio_url, is_pushed)
                        VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                        ON CONFLICT (id_original) DO UPDATE 
                        SET telegram_id = EXCLUDED.telegram_id,
                            time = EXCLUDED.time,
                            title = EXCLUDED.title,
                            content = EXCLUDED.content,
                            url = EXCLUDED.url,
                            mobile_url = EXCLUDED.mobile_url,
                            reading_num = EXCLUDED.reading_num,
                            audio_url = EXCLUDED.audio_url;
                    ''', (
                          telegram.get('id'),
                          telegram.get('time'),
                          telegram.get('title'),
                          telegram.get('content'),
                          telegram.get('url'),
                          telegram.get('mobileUrl'),
                          telegram.get('reading_num'),
                          telegram.get('id_original'),  # 这是真正的唯一标识符
                          telegram.get('audio_url'),
                          not web_config.get_wechat_push_enabled()  # 根据推送开关设置默认推送状态
                      ))
                    
                    if cursor.rowcount > 0:
                        updated_count += 1
                        
                except Exception as e:
                    print(f"保存电报数据时出错: {e}")
                    # 继续处理下一条数据
            
            conn.commit()
        except Exception as e:
            print(f"批量保存时发生错误: {e}")
            conn.rollback()
        finally:
            cursor.close()
            conn.close()
        
        return updated_count
        


    def get_latest_telegrams(self, limit=10):
        """获取最新的电报数据"""
        conn = self.get_connection()
        if not conn:
            return []

        try:
            cursor = conn.cursor()
            
            cursor.execute('''
                SELECT * FROM telegrams 
                ORDER BY time DESC 
                LIMIT %s
            ''', (limit,))
            
            results = cursor.fetchall()
            
            # 转换为字典格式
            columns = ['id', 'telegram_id', 'time', 'title', 'content', 'url', 
                      'mobile_url', 'reading_num', 'id_original', 'audio_url', 'created_at']
            return [dict(zip(columns, row)) for row in results]
        except Exception as e:
            print(f"获取最新电报数据时出错: {e}")
            return []
        finally:
            cursor.close()
            conn.close()

    def get_telegram_count(self):
        """获取电报总数"""
        conn = self.get_connection()
        if not conn:
            return 0

        try:
            cursor = conn.cursor()
            
            cursor.execute('SELECT COUNT(*) FROM telegrams')
            count = cursor.fetchone()[0]
            return count
        except Exception as e:
            print(f"获取电报总数时出错: {e}")
            return 0
        finally:
            cursor.close()
            conn.close()

    def get_unpushed_telegrams(self, limit=5):
        """获取未推送的电报数据"""
        conn = self.get_connection()
        if not conn:
            return []

        try:
            cursor = conn.cursor()
            
            cursor.execute('''
                SELECT * FROM telegrams 
                WHERE is_pushed = FALSE
                ORDER BY 
                    -- 优先尝试将time字段转换为时间戳进行排序，确保时间顺序准确
                    CASE 
                        WHEN time ~ '^\\d{4}-\\d{2}-\\d{2} \\d{2}:\\d{2}:\\d{2}$' THEN to_timestamp(time, 'YYYY-MM-DD HH24:MI:SS')
                        WHEN time ~ '^\\d{4}-\\d{2}-\\d{2}$' THEN to_timestamp(time, 'YYYY-MM-DD')
                        ELSE NULL
                    END DESC,
                    -- 如果时间转换失败，使用原始time字符串排序作为备选
                    time DESC 
                LIMIT %s
            ''', (limit,))
            
            results = cursor.fetchall()
            
            # 转换为字典格式
            columns = ['id', 'telegram_id', 'time', 'title', 'content', 'url', 
                      'mobile_url', 'reading_num', 'id_original', 'audio_url', 'created_at', 'is_pushed']
            return [dict(zip(columns, row)) for row in results]
        except Exception as e:
            print(f"获取未推送电报数据时出错: {e}")
            return []
        finally:
            cursor.close()
            conn.close()

    def mark_as_pushed(self, id_original):
        """标记电报为已推送"""
        import re
        import logging
        
        # 配置日志
        logger = logging.getLogger(__name__)
        logger.info(f"开始标记消息(id_original={id_original})为已推送")
        
        # 确保id_original是可处理的类型，避免将URL作为整数使用
        if isinstance(id_original, str) and id_original.startswith('http'):
            logger.warning(f"接收到URL而不是ID: {id_original}，尝试从中提取ID")
            # 从URL中提取文章ID，匹配/article/后面的数字
            match = re.search(r'article/(\d+)', id_original)
            if match:
                numeric_id = match.group(1)
                logger.info(f"从URL中提取到ID: {numeric_id}")
                # 尝试转换为整数，但如果失败则保持字符串形式
                try:
                    return self._mark_as_pushed_by_id(int(numeric_id), logger)
                except ValueError:
                    logger.warning(f"无法将提取的ID {numeric_id} 转换为整数")
                    return self._mark_as_pushed_by_id(numeric_id, logger)
            else:
                logger.warning(f"无法从URL中提取ID: {id_original}")
                # 如果无法提取ID，尝试直接使用url字段匹配
                return self._mark_as_pushed_by_url(id_original, logger)
        
        # 对于正常的ID，调用内部方法处理
        result = self._mark_as_pushed_by_id(id_original, logger)
        
        # 如果所有尝试都失败，执行最后的备选方案：记录未成功标记的消息
        if not result:
            logger.error(f"所有尝试都无法标记电报为已推送: {id_original}")
            self._log_push_failure(id_original, logger)
        else:
            logger.info(f"成功标记消息(id_original={id_original})为已推送")
            
        return result
        
    def _mark_as_pushed_by_id(self, id_original, logger):
        """内部方法：通过ID标记消息为已推送"""
        # 尝试多次执行更新操作，增加成功率
        retry_count = 3
        for attempt in range(retry_count):
            try:
                conn = self.get_connection()
                if not conn:
                    logger.error(f"第{attempt+1}次尝试：无法连接到数据库")
                    continue
                
                cursor = conn.cursor()
                
                # 策略1：尝试使用id_original更新（主要策略）
                try:
                    cursor.execute('''
                        UPDATE telegrams 
                        SET is_pushed = TRUE 
                        WHERE id_original = %s
                    ''', (id_original,))
                    
                    if cursor.rowcount > 0:
                        conn.commit()
                        logger.info(f"成功标记消息(id_original={id_original})为已推送")
                        cursor.close()
                        conn.close()
                        return True
                except Exception as e:
                    logger.warning(f"策略1执行出错: {e}")
                      
                # 策略2：如果是数字字符串，尝试转换为整数后更新
                if isinstance(id_original, str) and id_original.isdigit():
                    try:
                        numeric_id = int(id_original)
                        cursor.execute('''
                            UPDATE telegrams 
                            SET is_pushed = TRUE 
                            WHERE id_original = %s
                        ''', (numeric_id,))
                        
                        if cursor.rowcount > 0:
                            conn.commit()
                            logger.info(f"成功标记消息(数字ID={numeric_id})为已推送")
                            cursor.close()
                            conn.close()
                            return True
                    except Exception as e:
                        logger.warning(f"策略2执行出错: {e}")
                          
                # 策略3：尝试使用telegram_id字段
                try:
                    # 确保传入的是整数类型
                    if isinstance(id_original, str):
                        if id_original.isdigit():
                            telegram_id = int(id_original)
                        else:
                            # 如果不是数字字符串，跳过这个策略
                            raise ValueError("ID不是数字字符串")
                    else:
                        telegram_id = id_original
                        
                    cursor.execute('''
                        UPDATE telegrams 
                        SET is_pushed = TRUE 
                        WHERE telegram_id = %s
                    ''', (telegram_id,))
                    
                    if cursor.rowcount > 0:
                        conn.commit()
                        logger.info(f"成功标记消息(telegram_id={telegram_id})为已推送")
                        cursor.close()
                        conn.close()
                        return True
                except Exception as e:
                    logger.warning(f"策略3执行出错: {e}")
                          
                # 如果所有策略都失败，回滚事务
                conn.rollback()
                logger.warning(f"第{attempt+1}次尝试：所有策略都无法标记电报为已推送: {id_original}")
                
                cursor.close()
                conn.close()
                
                # 如果不是最后一次尝试，等待一段时间后重试
                if attempt < retry_count - 1:
                    time.sleep(1)  # 等待1秒后重试
                    
            except Exception as e:
                logger.error(f"第{attempt+1}次尝试标记消息为已推送时出错: {e}")
                if 'conn' in locals() and not conn.closed:
                    conn.rollback()
                    cursor.close()
                    conn.close()
                # 如果不是最后一次尝试，等待一段时间后重试
                if attempt < retry_count - 1:
                    time.sleep(1)
                    
        return False
        
    def _mark_as_pushed_by_url(self, url, logger):
        """内部方法：通过URL标记消息为已推送"""
        try:
            conn = self.get_connection()
            if not conn:
                logger.error("无法连接到数据库")
                return False
            
            cursor = conn.cursor()
            
            # 尝试直接使用url字段匹配
            cursor.execute('''
                UPDATE telegrams 
                SET is_pushed = TRUE 
                WHERE url = %s
            ''', (url,))
            
            if cursor.rowcount > 0:
                conn.commit()
                logger.info(f"成功标记消息(URL匹配={url})为已推送")
                cursor.close()
                conn.close()
                return True
            
            # 如果URL完全匹配失败，尝试部分匹配
            cursor.execute('''
                UPDATE telegrams 
                SET is_pushed = TRUE 
                WHERE url LIKE %s
            ''', (f'%{url}%',))
            
            if cursor.rowcount > 0:
                conn.commit()
                logger.info(f"成功标记消息(URL部分匹配={url})为已推送")
                cursor.close()
                conn.close()
                return True
            
            conn.rollback()
            logger.warning(f"无法通过URL标记电报为已推送: {url}")
            
            cursor.close()
            conn.close()
            
        except Exception as e:
            logger.error(f"通过URL标记消息为已推送时出错: {e}")
            if 'conn' in locals() and not conn.closed:
                conn.rollback()
                cursor.close()
                conn.close()
                
        return False
        
    def _log_push_failure(self, id_original, logger):
        """记录推送失败的消息"""
        try:
            conn = self.get_connection()
            if conn:
                cursor = conn.cursor()
                # 创建日志表（如果不存在）
                cursor.execute('''
                    CREATE TABLE IF NOT EXISTS push_failures (
                        id SERIAL PRIMARY KEY,
                        id_original TEXT,
                        failure_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                    )
                ''')
                # 插入失败记录
                cursor.execute('''
                    INSERT INTO push_failures (id_original) VALUES (%s)
                ''', (str(id_original),))
                conn.commit()
                cursor.close()
                conn.close()
                logger.info(f"已记录推送标记失败: {id_original}")
        except Exception as e:
            logger.error(f"记录推送失败时出错: {e}")
    
    def get_telegrams_by_date_range(self, start_date, end_date):
        """根据日期范围获取电报"""
        conn = self.get_connection()
        if not conn:
            return []

        try:
            cursor = conn.cursor()
            
            cursor.execute('''
                SELECT * FROM telegrams 
                WHERE time BETWEEN %s AND %s
                ORDER BY time DESC
            ''', (start_date, end_date))
            
            results = cursor.fetchall()
            
            # 转换为字典格式
            columns = ['id', 'telegram_id', 'time', 'title', 'content', 'url', 
                      'mobile_url', 'reading_num', 'id_original', 'audio_url', 'created_at']
            return [dict(zip(columns, row)) for row in results]
        except Exception as e:
            print(f"根据日期范围获取电报时出错: {e}")
            return []
        finally:
            cursor.close()
            conn.close()